//OMP c++ wrapper
//$Id: easy_omp.h 719 2012-06-06 17:30:15Z Oleg.Bulychov@gmail.com $
#pragma once

//#include <deque>
#include <omp.h>
#include "src/heo/include/fast_vector.h"
#include "src/heo/include/reflection.h"
#include "src/heo/include/transmit.h"
#include "src/heo/include/sys_api.h"

//_______________________________________________________________________________
//_______________________________________________________________________________
namespace omp
{
//_______________________________________________________________________________
typedef size_t opid_t;
//_______________________________________________________________________________
class mutex
{
    DISALLOW_COPY_AND_ASSIGN(mutex);
    omp_lock_t handle_;
public:
    mutex()
    {
        omp_init_lock(&handle_);
    }
    ~mutex()
    {
        omp_destroy_lock(&handle_);
    }

    class scoped_lock
    {
        DISALLOW_COPY_AND_ASSIGN(scoped_lock);
        mutex* m_;
    public:
        scoped_lock(): m_(NULL)
        {}

        explicit scoped_lock(mutex& m): m_(&m)
        {
            omp_set_lock(&m.handle_);
        }

        ~scoped_lock()
        {
            if (m_)
                omp_unset_lock(&m_->handle_);
        }

        bool try_acquire(mutex& m)
        {
            ASSERT_(!m_);
            if (omp_test_lock(&m.handle_))
            {
                m_ = &m;
                return true;
            }
            return false;
        }
    };
};
//_______________________________________________________________________________
/*
template<typename T = int>
class atomic
{
    volatile T value_;
    DISALLOW_COPY_AND_ASSIGN(atomic);
public:
    atomic(T v = 0): value_(v)
    {}

    atomic<T>& operator ++()
    {
#pragma omp atomic
        ++value_;
        return *this;
    }

    atomic<T>& operator --()
    {
#pragma omp atomic
        --value_;
        return *this;
    }

    operator volatile T()
    {
        return value_;
    }
};
*/
//_______________________________________________________________________________
static const unsigned int OMP_ISEND_DEFAULT  = 0;
static const unsigned int OMP_ISEND_BUFFERED = 1;
//_______________________________________________________________________________
namespace hidden
{
//_______________________________________________________________________________
struct pending_data
{
    union
    {
        byte_t*   pv_;
        uint64_t  dummy_;
    };
    size_t        cb_;
    bool          need_free_pv_;
};
//_______________________________________________________________________________
//_______________________________________________________________________________
//_______________________________________________________________________________
template < class M = mutex, class P = fast_vector<pending_data> >
class thread_data_buffer
{
    DISALLOW_COPY_AND_ASSIGN(thread_data_buffer);
public:
    typedef M                               TMutex;
    typedef P                               TPipe;
    typedef typename TMutex::scoped_lock    TLock;
//.............................................................................
    TPipe*           pipe_;
    volatile opid_t* head_op_id_;
    volatile opid_t* tail_op_id_;
    TMutex           mutex_;
    int              num_threads_;
    volatile int     data_avail_;
//.............................................................................
    thread_data_buffer():
        pipe_(NULL),
        head_op_id_(NULL),
        tail_op_id_(NULL),
        mutex_(),
        num_threads_(0),
        data_avail_(0)
    {}
//.............................................................................
    ~thread_data_buffer()
    {
        free_Pipe_Data();
        delete[] pipe_;
        delete[] head_op_id_;
        delete[] tail_op_id_;
    }
//.............................................................................
    void init(int num_threads)
    {
        num_threads_ = num_threads;
        pipe_ = new TPipe[num_threads];
        head_op_id_ = new opid_t[num_threads];
        tail_op_id_ = new opid_t[num_threads];
        for (int i = 0; i < num_threads_; ++i)
        {
            head_op_id_[i] = 0;
            tail_op_id_[i] = 0;
        }
    }
//.............................................................................
    void clear_Data()
    {
        free_Pipe_Data();
        for (int i = 0; i < num_threads_; ++i)
        {
            pipe_[i].clear();
            head_op_id_[i] = 0;
            tail_op_id_[i] = 0;
        }
        data_avail_ = 0;
    }
//.............................................................................
    void free_Pipe_Data()
    {
        for (int i = 0; i < num_threads_; ++i)
        {
            TPipe& pipe = pipe_[i];
            for (size_t j = 0; j < pipe.size(); ++j)
            {
                if (pipe[j].need_free_pv_)
                    free(pipe[j].pv_);
            }
        }
    }
//.............................................................................
};
//_______________________________________________________________________________
class thread_buffer_alloc
{
public:
    static void* Alloc(size_t cb)
    {
        return malloc(cb);
    }
    static void Free(void * pv, size_t)
    {
        return free(pv);
    }
};
//_______________________________________________________________________________
template<class A = thread_buffer_alloc>
class thread_buffer_list
{
    DISALLOW_COPY_AND_ASSIGN(thread_buffer_list);
public:
    typedef thread_data_buffer<>   TList;
    typedef TList::TLock           TLock;
    typedef TList::TPipe           TPipe;
    typedef thread_buffer_alloc    TAlloc;
private:
    TList*  list_;
    int num_threads_;
public:
    thread_buffer_list(): list_(NULL), num_threads_(0)
    {}
//.............................................................................
    ~thread_buffer_list()
    {
        delete[] list_;
    }
//.............................................................................
    void init(int num_threads)
    {//:NOTE: not thread safe!
        num_threads_ = num_threads;
        list_ = new TList[num_threads];
        for (int i = 0; i < num_threads; ++i)
            list_[i].init(num_threads);
    }
//.............................................................................
    void clear_Data(int rank)
    {
        ASSERT_(rank < num_threads_);

        TLock lock(list_[rank].mutex_);
        list_[rank].clear_Data();
    }
//.............................................................................
    opid_t isend(int rank, int dst, byte_t* pv, size_t cb, unsigned int flags)
    {
        ASSERT_(rank < num_threads_);
        ASSERT_(dst < num_threads_);

        {
            TLock lock(list_[dst].mutex_);
            {
                TPipe& pipe = list_[dst].pipe_[rank];
                pending_data data;
                data.cb_ = cb;
                if (cb <= sizeof(data.dummy_))
                {
                    memcpy(&data.dummy_, pv, cb);
                    data.need_free_pv_ = false;
                }
                else if (flags & OMP_ISEND_BUFFERED)
                {
                    byte_t* buffer = (byte_t*)TAlloc::Alloc(cb);
                    memcpy(buffer, pv, cb);
                    data.pv_= buffer;
                    data.need_free_pv_ = true;
                }
                else
                {
                    data.pv_ = pv;
                    data.need_free_pv_ = false;
                }
                pipe.push_back(data);
                ++list_[dst].data_avail_;
                return ++list_[dst].head_op_id_[rank];
            }
        }
    }
//.............................................................................
    void send(int rank, int dst, byte_t* pv, size_t cb)
    {
        opid_t op_id = isend(rank, dst, pv, cb, OMP_ISEND_DEFAULT);

        for(;;)
        {
            if (is_Op_Complete(rank, dst, op_id))
                break;
            sys_api::Yield_Thread();
        }
        //while(list_[dst].data_avail_)
        //{
        //    sys_api::Yield_Thread();
        //    TLock lock;
        //    if (lock.try_acquire(list_[dst].mutex_))
        //        if (list_[dst].pipe_[rank].empty())
        //            break;
        //}
    }
//.............................................................................
    bool is_Op_Complete(int rank, int dst, opid_t id)
    {
        //??TLock lock(list_[dst].mutex_);//:TODO review: avoid it?
        {
            opid_t tail_id = list_[dst].tail_op_id_[rank];
            return id <= tail_id;
        }
    }
//.............................................................................
    bool is_Data_Available(int rank, int src)
    {
        ASSERT_(rank < num_threads_);
        ASSERT_(src < num_threads_);

        if (!list_[rank].data_avail_)
            return false;
        TLock lock;
        if (!lock.try_acquire(list_[rank].mutex_))
            return false;//:dangerous
        return !list_[rank].pipe_[src].empty();
    }
//.............................................................................
    bool irecv(int rank, int src, byte_t* pv, size_t cb)
    {
        ASSERT_(rank < num_threads_);
        ASSERT_(src < num_threads_);
        if (!list_[rank].data_avail_)
            return false;
        TLock lock;
        if (!lock.try_acquire(list_[rank].mutex_))
            return false;
        TPipe& pipe = list_[rank].pipe_[src];
        if (pipe.empty())
            return false;
        {
            pending_data& data = pipe.front();
            VERIFY_(cb == data.cb_);
            if (cb <= sizeof(data.dummy_))
                memcpy(pv, &data.dummy_, cb);
            else
                memcpy(pv, data.pv_, cb);
            if (data.need_free_pv_)
                TAlloc::Free(data.pv_, cb);
        }
        pipe.erase(pipe.begin(), pipe.begin() + 1);//pop_front();
        ++list_[rank].tail_op_id_[src];
        --list_[rank].data_avail_;
        //std::cout << "irecv " << cb << std::endl;
        return true;
    }
//.............................................................................
    void recv(int rank, int src, byte_t* pv, size_t cb)
    {
        while(!irecv(rank, src, pv, cb))
        {
            sys_api::Yield_Thread();
        }
    }
//.............................................................................
//.............................................................................
//.............................................................................
};
//_______________________________________________________________________________
//_______________________________________________________________________________
//_______________________________________________________________________________
//_______________________________________________________________________________
template<class M>
class omp_visitor:
    public meta::visitor<M>
{
public:
    typedef meta::visitor< M > TVisitor;
    using TVisitor::self;
//.............................................................................
    template<typename T, typename D>
    void visit(T& t, char const* szName, D const& d)
    {
        visit2(t, meta::int2type<meta::is_integral_type<T>::result>());
    }
    template<typename T>
    void visit2(T& t, meta::int2type<0>)
    {
        transmit_handler<T, M>::visit(t, *self());
    }
    template<typename T>
    void visit2(T& t, meta::int2type<1>)
    {
        self()->invoke(&t, 1);
    }
//.............................................................................
    template<typename T, typename D>
    void visit_Array(T* pt, size_t N, char const* szName, D const& d)
    {
        visit_Array2(pt, N, szName, d, meta::int2type<meta::is_integral_type<T>::result>());
    }
    template<typename T, typename D>
    void visit_Array2(T* pt, size_t N, char const* szName, D const& d, meta::int2type<0> it)
    {
        TVisitor::visit_Array(pt, N, szName, d);
    }
    template<typename T, typename D>
    void visit_Array2(T* pt, size_t N, char const* szName, D const& d, meta::int2type<1>)
    {
        self()->invoke(pt, N);
    }
//.............................................................................
};
//_______________________________________________________________________________
}//hidden
//_______________________________________________________________________________
//_______________________________________________________________________________
//_______________________________________________________________________________
//_______________________________________________________________________________
class communicator
{
    typedef hidden::thread_buffer_alloc     TAlloc;

    hidden::thread_buffer_list<TAlloc>&     itc_;
    int                                     num_threads_;
    int                                     rank_;
public:
    communicator(hidden::thread_buffer_list<TAlloc>& itc, int num_threads, int rank):
        itc_(itc), num_threads_(num_threads), rank_(rank)
    {}
//.............................................................................
    template<class T>
    static void run(int num_threads, T& t)
    {//:NOTE: call from main thread!
        hidden::thread_buffer_list<TAlloc> itc;
        itc.init(num_threads);
#pragma omp parallel num_threads(num_threads)
        {
            int rank = omp_get_thread_num();//:The master thread of the team returns a value of 0
            communicator tg(itc, num_threads, rank);
            t(tg);
        }
    }
//.............................................................................
    template<class T, class Shared>
    static void run(int num_threads, T& t, Shared& shared)
    {//:NOTE: call from main thread!
        hidden::thread_buffer_list<TAlloc> itc;
        itc.init(num_threads);
#pragma omp parallel num_threads(num_threads)
        {
            int rank = omp_get_thread_num();//:The master thread of the team returns a value of 0
            communicator tg(itc, num_threads, rank);
            t(tg, shared);
        }
    }
//.............................................................................
    int size() const
    {
        return num_threads_;
    }
//.............................................................................
    int rank() const
    {
        return rank_;
    }
//.............................................................................
    void barrier()
    {
#pragma omp barrier
    }
//.............................................................................
    void send(int dst, byte_t* pv, size_t cb)
    {
        itc_.send(rank_, dst, pv, cb);
    }
//.............................................................................
    opid_t isend(int dst, byte_t* pv, size_t cb, unsigned int flags)
    {
        return itc_.isend(rank_, dst, pv, cb, flags);
    }
//.............................................................................
    bool is_Op_Complete(int dst, opid_t id)
    {
        return itc_.is_Op_Complete(rank_, dst, id);
    }
//.............................................................................
    bool is_Data_Available(int src)
    {
        return itc_.is_Data_Available(rank_, src);
    }
//.............................................................................
    void recv(int src, byte_t* pv, size_t cb)
    {
        itc_.recv(rank_, src, pv, cb);
    }
//.............................................................................
    template<typename T>
    void all_gather(T* in, int count, T* out) const
    {//:out must be shared!
        T* dst = out;
        dst += count * rank_;
        memcpy(dst, in, count * sizeof(T));
#pragma omp barrier
    }
//.............................................................................
//.............................................................................
//.............................................................................
    void clear_Data()
    {
        itc_.clear_Data(rank_);
    }
};
//_______________________________________________________________________________
class send:
    public hidden::omp_visitor<send>
{
    int dst_;
    communicator& comm_;
public:
    enum { transmit_tag = SEND_TAG };
    bool buffered_;
public:
    send(int dest, communicator& comm): dst_(dest),
        comm_(comm)
    {}
//.............................................................................
    template<typename T>
    send& operator <<(T& t)
    {
        visit_T(t);
        return *this;
    }
//.............................................................................
    template<typename T>
    void invoke(T* pt, size_t count)
    {
        COMPILETIME_CHECK(meta::is_integral_type<T>::result, bad_value);
        comm_.send(dst_, (byte_t*)pt, count * sizeof(T));
    }
};
//_______________________________________________________________________________
class isend:
    public hidden::omp_visitor<isend>
{
    int dst_;
    opid_t op_id_;
    communicator& comm_;
public:
    enum { transmit_tag = SEND_TAG };
    bool buffered_;
public:
    isend(int dest, communicator& comm): dst_(dest),
        op_id_(0),
        comm_(comm),
        buffered_(false)
    {}
//.............................................................................
    template<typename T>
    isend& operator <<(T& t)
    {
        visit_T(t);
        return *this;
    }
//.............................................................................
    template<typename T>
    void invoke(T* pt, size_t count)
    {
        COMPILETIME_CHECK(meta::is_integral_type<T>::result, bad_value);
        op_id_ = comm_.isend(dst_, (byte_t*)pt, count * sizeof(T), buffered_ ? OMP_ISEND_BUFFERED : OMP_ISEND_DEFAULT);
    }
//.............................................................................
    opid_t get_Last_Op_Id() const
    {
        return op_id_;
    }
};
//_______________________________________________________________________________
//TODO kill?
class ibsend:
    public hidden::omp_visitor<ibsend>
{
    int dst_;
    opid_t op_id_;
    communicator& comm_;
public:
    enum { transmit_tag = SEND_TAG };
public:
    ibsend(int dest, communicator& comm): dst_(dest),
        op_id_(0),
        comm_(comm)
    {}
//.............................................................................
    template<typename T>
    ibsend& operator <<(T& t)
    {
        visit_T(t);
        return *this;
    }
//.............................................................................
    template<typename T>
    void invoke(T* pt, size_t count)
    {
        COMPILETIME_CHECK(meta::is_integral_type<T>::result, bad_value);
        op_id_ = comm_.isend(dst_, (byte_t*)pt, count * sizeof(T), OMP_ISEND_BUFFERED);
    }
//.............................................................................
    opid_t get_Last_Op_Id() const
    {
        return op_id_;
    }
};
//_______________________________________________________________________________
class recv:
    public hidden::omp_visitor<recv>
{
    int src_;
    communicator& comm_;
public:
    enum { transmit_tag = RECV_TAG };
public:
    recv(int src, communicator& comm): src_(src),
        comm_(comm)
    {}
//.............................................................................
    template<typename T>
    recv& operator >>(T& t)
    {
        visit_T(t);
        return *this;
    }
//.............................................................................
    template<typename T>
    void invoke(T* pt, size_t count)
    {
        COMPILETIME_CHECK(meta::is_integral_type<T>::result, bad_value);
        comm_.recv(src_, (byte_t*)pt, count * sizeof(T));
    }
//.............................................................................
    bool is_Data_Available()
    {
        return comm_.is_Data_Available(src_);
    }
};
//_______________________________________________________________________________
class bcast:
    public hidden::omp_visitor<bcast>
{
    int root_;
    communicator& comm_;
public:
    enum { transmit_tag = BCAST_TAG };
public:
    bcast(int root, communicator& comm): root_(root),
        comm_(comm)
    {}
//.............................................................................
    int get_Root() const { return root_; }
//.............................................................................
    bool is_Root()
    {
        return comm_.rank() == root_;
    }
//.............................................................................
    template<typename T>
    bcast& operator >>(T& t)
    {
        visit_T(t);
        return *this;
    }
//.............................................................................
    template<typename T>
    bcast& operator <<(T& t)
    {
        visit_T(t);
        return *this;
    }
//.............................................................................
    template<typename T>
    void invoke(T* pt, size_t count)
    {
        COMPILETIME_CHECK(meta::is_integral_type<T>::result, bad_value);
        if (is_Root())
        {
            for (int i = 1; i < comm_.size(); ++i)
                comm_.send(i, (byte_t*)pt, count * sizeof(T));
        }
        else
        {
            comm_.recv(0, (byte_t*)pt, count * sizeof(T));
        }
    }
};
//_______________________________________________________________________________
class dummy_thread_handler
{
public:
    static void enter_Thread(int nRank) {}
    static void leave_Thread(int nRank) {}
};
//_______________________________________________________________________________
template<class TSolutionPool, class TStatistics, class TSolver,
    class ThreadHandler = dummy_thread_handler>
struct solver_fork
{
    TSolutionPool&  pool_;
    TStatistics&    stat_;
    TSolver&        s_;
//.............................................................................
    explicit solver_fork(TSolutionPool& pool, TStatistics& stat, TSolver& s):
        pool_(pool),
        stat_(stat),
        s_(s)
    {}
//.............................................................................
    void operator()(omp::communicator& comm)
    {
        int rank = comm.rank();
        stat_.enter_Thread(rank);
        if (s_.get_Rank() == rank)
        {
            s_.solve_OMP(pool_, stat_, comm);
        }
        else
        {
            ThreadHandler::enter_Thread(rank);
            {
                TSolver s;
                s.init(s_.get_Environment(), rank, rank);
                xml::in(s_.get_Environment()->property_tree_) >> s;
                s.solve_OMP(pool_, stat_, comm);
            }
            ThreadHandler::leave_Thread(rank);
        }
        stat_.leave_Thread(rank);
    }
//.............................................................................
    template<class S>
    void operator()(omp::communicator& comm, S& shared)
    {
        int rank = comm.rank();
        stat_.enter_Thread(rank);
        if (s_.get_Rank() == rank)
        {
            s_.solve_OMP(pool_, stat_, comm, shared);
        }
        else
        {
            ThreadHandler::enter_Thread(rank);
            {
                TSolver s;
                s.init(s_.get_Environment(), rank, rank);
                xml::in(s_.get_Environment()->property_tree_) >> s;
                s.solve_OMP(pool_, stat_, comm, shared);
            }
            ThreadHandler::leave_Thread(rank);
        }
        stat_.leave_Thread(rank);
    }
};
//_______________________________________________________________________________

}//omp
//__________________________________________________________________________________
//__________________________________________________________________________________
//EOF!
